from transformers import AutoModel, AutoConfig
import yaml

# 加载预训练的模型
# model_opt = OPTModel.from_pretrained("/home/lthpc/nvmessd/wangjiaqian/simulator-by-alpa/model/opt_125m")
# model_qwen2 = AutoModel.from_pretrained("/home/lthpc/nvmessd/wangjiaqian/simulator-by-alpa/model/qwen2_0.5b")
model_qwen3_8b = AutoModel.from_pretrained("/home/lionrock-g4/wangjiaqian/data/wjq_data/model/Qwen3-32B")
# model_qwen3_8b = AutoModel.from_pretrained("Qwen/Qwen3-8B")
model_name = "Qwen/Qwen3-8B"
model_config_path = "/pfs/data/wangjiaqian/llm_simulator_by_alpa_layer/config/model_config_qwen3_8b.yaml"

def get_model_structure(model):
    model_structure = {"layers": []}

    for name, layer in model.named_modules():
        try:
            # 打印每一层的名称和类型
            print(f"Layer name: {name}, Layer type: {type(layer).__name__}")

            # 将所有层的信息添加到结构中
            model_structure["layers"].append({
                "name": name,
                "type": type(layer).__name__,
                "params": {k: v for k, v in layer.named_parameters()}
            })

        except Exception as e:
            print(f"忽略未识别的层 {name}: {e}")

    return model_structure


def get_model_parameter_in_yaml(model_name, model_config_path):
    config = AutoConfig.from_pretrained(model_name)
    # 提取超参数
    model_config = {
        "model_name": model_name,
        "vocab_size": config.vocab_size,
        "hidden_dim": config.hidden_size,
        "num_layers": config.num_hidden_layers,
        "num_heads": config.num_attention_heads,
        "mlp_dim": config.intermediate_size,  # MLP 层隐藏维度
        "max_seq_length": config.max_position_embeddings,
    }

    with open(model_config_path, "w") as file:
        yaml.dump(model_config, file, default_flow_style=False, allow_unicode=True)


if __name__ == "__main__":

    print("\n=====  Model Structure =====")
    qwen2_structure = get_model_structure(model_qwen3_8b)
    print(f"Qwen3-8B Model has {len(qwen2_structure['layers'])} layers.")

    print("\n===== Saving  Model Config to YAML =====")
    get_model_parameter_in_yaml(model_name, model_config_path)
    print(f"模型参数已保存到: {model_config_path}")





